"""类型替换Pass"""

from typing import Optional, Dict, Any
from mcpy.ir import IRNode, IRModule, IRFunctionNode
from mcpy.ir.declarations.method import IRMethodNode
from mcpy.ir.statements.assignment import IRAssign
from mcpy.ir.statements import IRReturn
from mcpy.ir.expressions.lambda_expr import IRLambda
from mcpy.ir.expressions import IRCall, IRConditional, IRName, IRBinOp, IRSubscript
from mcpy.symbol_meta import TypeInfo, TypeKind
from mcpy.passes.base import BasePass, PassStage, PassContext
from mcpy.type_system import (
    get_type_includes,
    get_cpp_type,
)


# 需要进行类型转换的节点类型
#
# 设计原则：
# 1. 只转换在代码生成阶段需要 C++ 类型信息的节点
# 2. 避免转换纯结构性节点（如 class, module, import 等）
# 3. 优先转换有具体值或需要类型推断的节点
CONVERTIBLE_NODE_TYPES = {
    # 声明节点 - 需要类型信息用于代码生成
    "property",  # 属性声明：int x = 5 -> int x = 5;
    "parameter",  # 函数参数：def func(x: int) -> void func(int x)
    "function",  # 函数声明：def func() -> int -> int func()
    "method",  # 方法声明：def method(self) -> str -> std::string method()
    # 表达式节点 - 需要类型信息用于类型推断和代码生成
    "call",  # 函数调用：func(args) -> 需要返回类型用于后续推断
    "name",  # 变量引用：x -> 需要变量类型用于代码生成
    "constant",  # 常量：42, "hello" -> 需要字面量类型
    "attribute",  # 属性访问：obj.attr -> 需要属性类型
    "subscript",  # 下标访问：arr[0] -> 需要元素类型
    "binop",  # 二元运算：a + b -> 需要结果类型
    "bin_op",  # 二元运算（IR节点实际使用 bin_op）
    "unaryop",  # 一元运算：-x -> 需要结果类型
    "unary_op",  # 一元运算（IR节点实际使用 unary_op）
    "compare",  # 比较运算：a > b -> 需要结果类型（通常是 bool）
    "conditional",  # 三元表达式：x if cond else y -> 需要结果类型
    "list",  # 列表字面量：[1, 2, 3] -> mc::array<int>
    "dict",  # 字典字面量：{"a": 1} -> mc::dict 或 std::map
    "tuple",  # 元组字面量：(1, 2) -> std::tuple<int, int>
    "set",  # 集合字面量：{1, 2} -> std::set<int>
    "list_comp",  # 列表推导式：[x for x in nums] -> mc::array<T>
    "set_comp",  # 集合推导式：{x for x in nums} -> std::set<T>
    "dict_comp",  # 字典推导式：{k: v for x in items} -> mc::dict 或 std::map
    "lambda",  # Lambda 表达式：lambda x: x + 1 -> 函数对象
    "await",  # Await 表达式：await future<T> -> T
    # 语句节点 - 部分需要类型信息
    "assign",  # 赋值语句：x = value -> 需要变量类型用于声明
    "return",  # 返回语句：return value -> 需要返回类型用于类型检查
    # 不包含的节点类型（说明为什么不需要转换）：
    # "class"       - 类定义本身不需要 C++ 类型，只是结构定义
    # "module"      - 模块不需要 C++ 类型
    # "import"      - 导入语句不需要 C++ 类型
    # "comment"     - 注释不需要 C++ 类型
    # "annotation"  - 注解节点不需要 C++ 类型（注解内容已被处理）
    # "if"/"while"  - 控制流语句本身不需要类型，只有其条件表达式需要
    # "enum"        - 枚举定义本身不需要 C++ 类型
}


class TypeReplacementPass(BasePass):
    """类型替换Pass

    职责：
    1. 将已推导的TypeInfo转换为C++类型
    2. 收集所需的头文件
    """

    name = "type_replacement"
    stage = PassStage.POST_ANALYSIS

    def __init__(self, config: Optional[Dict[str, Any]] = None):
        super().__init__()
        self.config = config or {}

        # 允许通过配置自定义需要转换的节点类型
        custom_types = self.config.get("convertible_node_types")
        if custom_types:
            self.convertible_node_types = set(custom_types)
        else:
            self.convertible_node_types = CONVERTIBLE_NODE_TYPES

    def run(self, ir_module: IRModule, context: PassContext) -> IRModule:
        """运行类型替换Pass

        使用后序遍历（先处理子节点，再处理父节点），
        以确保嵌套表达式（如 d["a"]["b"]）能正确获取子节点的 cpp_type
        """
        # 后序遍历：先处理子节点，再处理当前节点
        self._postorder_convert(ir_module, ir_module)
        return ir_module

    def _postorder_convert(self, node: IRNode, ir_module: IRModule):
        """后序遍历并转换节点类型"""
        for child in node.get_child_nodes():
            self._postorder_convert(child, ir_module)

        self._analyze_lambda_capture_strategy(node, ir_module)
        self._analyze_nested_function_capture_strategy(node, ir_module)
        if self._should_convert_node(node):
            self._convert_node_type(node, ir_module)
        self._convert_loop_var_types(node, ir_module)
        self._fix_assignment_target_type(node, ir_module)

        if isinstance(node, IRFunctionNode):
            # 函数级别后处理
            # 1. 标记需要 shared_ptr 的变量
            self._mark_shared_ptr_variables(node, ir_module)

            # 2. 更新栈 lambda 引用
            for child in node.walk():
                self._fix_stack_lambda_references(child, ir_module)

            # 3. 转换异步状态变量类型
            if node.is_async:
                self._convert_async_state_vars(node)

    def _should_convert_node(self, node: IRNode) -> bool:
        """判断节点是否需要进行类型转换

        Args:
            node: IR节点

        Returns:
            bool: 是否需要转换
        """
        # 检查节点类型是否在可转换列表中
        if isinstance(node, IRFunctionNode):
            # 函数节点及其子类都需要转换
            pass
        elif node.node_type not in self.convertible_node_types:
            return False

        if isinstance(node, IRLambda):
            return False

        # 检查是否已经转换过
        existing_cpp_type = node.get_extension_param("cpp_type")
        if existing_cpp_type:
            return False  # 已经有 C++ 类型，跳过

        # 特殊节点（如 bin_op）即使没有有效的 inferred_type 也需要处理
        if node.node_type in ("binop", "bin_op", "subscript", "conditional"):
            return True

        # 检查节点是否有推导的类型信息
        type_info = node.get_extension_param("inferred_type")
        if not type_info:
            return False

        return True

    def _convert_node_type(self, node: IRNode, ir_module: IRModule):
        """转换节点的类型信息

        Returns:
            bool: 是否成功转换了类型
        """
        # 先处理特殊节点（可能不需要 inferred_type）
        # 这些节点即使没有有效的 inferred_type 也需要处理

        # 特殊处理: list[Any] 应该生成 mc::variants
        # list[Any] 是泛型类型，其类型参数是 Any
        # 注意：这个处理需要递归处理嵌套类型，例如 list[list[Any]]
        def convert_list_any_to_variants(type_info):
            """递归转换 list[Any] 及其嵌套类型为 C++ 类型"""
            if not type_info or type_info.name != "list":
                return None

            # 检查是否是 list[Any]
            if type_info.args and len(type_info.args) == 1:
                arg_type = type_info.args[0]
                if arg_type.kind == TypeKind.ANY:
                    # list[Any] → mc::variants
                    return "mc::variants"
                elif arg_type.name == "list":
                    # 嵌套 list，递归处理
                    inner_cpp = convert_list_any_to_variants(arg_type)
                    if inner_cpp:
                        # list[list[Any]] → mc::array<mc::variants>
                        return f"mc::array<{inner_cpp}>"
            return None

        node_type_info = node.get_extension_param("inferred_type")
        if node_type_info and node_type_info.name == "list":
            cpp_type_result = convert_list_any_to_variants(node_type_info)
            if cpp_type_result:
                node.set_extension_param("cpp_type", cpp_type_result)
                ir_module.require_include("<mc/variant.h>")
                ir_module.require_include("<mc/array.h>")
                return
            # 无类型参数的 list（通过 default_type 配置已处理）不需要特殊处理

        # 特殊处理: dict 的下标访问
        # 由于 mc::dict 的 operator[] 总是返回 mc::variant,
        # 即使 Python 类型推导为 dict[K,V] 的值类型,C++ 类型也应该是 mc::variant
        # 同时，mc::variant 也可以表现为 dict（通过 .as_dict()），所以 variant["key"] 也返回 variant
        if node.node_type == "subscript":
            if isinstance(node, IRSubscript):
                # 检查容器是否是 dict 类型（通过 inferred_type）
                container_type = node.value.get_extension_param("inferred_type")
                is_dict = container_type and container_type.name == "dict"

                # 检查容器是否是 mc::variant（通过 cpp_type）
                # 这种情况发生在嵌套访问：d["a"]["b"]，其中 d["a"] 返回 variant
                container_cpp_type = node.value.get_extension_param("cpp_type")
                is_variant = container_cpp_type == "mc::variant"

                if is_dict or is_variant:
                    # dict 下标访问或 variant 下标访问在 C++ 中总是返回 mc::variant
                    node.set_extension_param("cpp_type", "mc::variant")
                    ir_module.require_include("<mc/variant.h>")
                    return

        # 特殊处理: 函数调用表达式
        # 函数调用的返回类型应该是被调用函数的返回类型，而不是 any
        if node.node_type == "call":
            if isinstance(node, IRCall) and isinstance(node.func, IRName):
                # 查找被调用的函数定义
                func_name = node.func.id
                for module_node in ir_module.nodes:
                    if (
                        isinstance(module_node, IRFunctionNode)
                        and module_node.name == func_name
                    ):
                        # 获取函数的返回类型（已经被 type_replacement 处理过）
                        func_return_cpp_type = module_node.get_cpp_return_type()
                        if func_return_cpp_type:
                            # 使用函数的返回类型作为调用表达式的类型
                            node.set_extension_param("cpp_type", func_return_cpp_type)
                            return

        # 特殊处理: 三元表达式（条件表达式）
        # 如果 body 或 orelse 任一分支是 any 类型，结果应该是 mc::variant
        if node.node_type == "conditional":
            if isinstance(node, IRConditional):
                # 检查 body 分支是否是 any
                body_type = node.body.get_extension_param("inferred_type")
                body_cpp = node.body.get_extension_param("cpp_type")
                body_is_any = (
                    body_type and body_type.kind == TypeKind.ANY
                ) or body_cpp == "mc::variant"

                # 检查 orelse 分支是否是 any
                orelse_type = node.orelse.get_extension_param("inferred_type")
                orelse_cpp = node.orelse.get_extension_param("cpp_type")
                orelse_is_any = (
                    orelse_type and orelse_type.kind == TypeKind.ANY
                ) or orelse_cpp == "mc::variant"

                # 如果任一分支是 any 类型，结果也是 mc::variant
                if body_is_any or orelse_is_any:
                    node.set_extension_param("cpp_type", "mc::variant")
                    ir_module.require_include("<mc/variant.h>")
                    return

        # 特殊处理: 二元运算表达式
        # 如果操作数包含 any 类型，结果类型应该是 mc::variant
        # 这是因为 variant 的运算符重载返回 variant
        if node.node_type in ("binop", "bin_op"):
            if isinstance(node, IRBinOp):
                # 检查左操作数是否是 any 类型
                left_type = node.left.get_extension_param("inferred_type")
                left_cpp = node.left.get_extension_param("cpp_type")
                left_is_any = (
                    left_type and left_type.kind == TypeKind.ANY
                ) or left_cpp == "mc::variant"

                # 检查右操作数是否是 any 类型
                right_type = node.right.get_extension_param("inferred_type")
                right_cpp = node.right.get_extension_param("cpp_type")
                right_is_any = (
                    right_type and right_type.kind == TypeKind.ANY
                ) or right_cpp == "mc::variant"

                # 如果任一操作数是 any 类型，结果也是 mc::variant
                # 因为 mc::variant 的所有运算符重载都返回 variant
                if left_is_any or right_is_any:
                    node.set_extension_param("cpp_type", "mc::variant")
                    ir_module.require_include("<mc/variant.h>")
                    return

        # 检查 type_info 是否有效（之前已经提前处理了特殊节点）
        type_info: TypeInfo = node.get_extension_param("inferred_type")
        if not type_info:
            # 对于函数节点，即使没有 type_info，也需要处理返回类型
            if isinstance(node, IRFunctionNode) and node.return_type:
                # 从 node.return_type 解析返回类型
                from mcpy.type_system import parse_type

                return_type_to_convert = parse_type(node.return_type, node)
                if return_type_to_convert:
                    # 特殊处理：None 作为函数返回类型时映射为 void
                    if return_type_to_convert.kind == TypeKind.NONE:
                        return_cpp_type = "void"
                    else:
                        return_cpp_type = get_cpp_type(return_type_to_convert, node)
                    if return_cpp_type:
                        node.set_extension_param("cpp_return_type", return_cpp_type)
            return

        # 检查是否已经设置了 cpp_type
        existing_cpp_type = node.get_extension_param("cpp_type")

        # 从符号配置获取 C++ 类型
        cpp_type = get_cpp_type(type_info, node)
        if cpp_type:
            # 处理未解析的模板参数占位（如 mc::future<T> → 回退 mc::variant）
            try:
                if type_info.name in ("future", "Future") and type_info.args:
                    arg0 = type_info.args[0]
                    if getattr(arg0, "kind", None) == TypeKind.TYPE_PARAM:
                        cpp_type = "mc::future<mc::variant>"
                        ir_module.require_include("<mc/future.h>")
            except Exception:
                pass
            # 如果已经有 cpp_type 且不是 mc::variant，保持原值（栈闭包场景）
            if (
                existing_cpp_type
                and existing_cpp_type != "mc::variant"
                and (
                    existing_cpp_type.startswith("__lambda_")
                    or existing_cpp_type.startswith("__nested_func_")
                )
            ):
                # 栈闭包（lambda 或嵌套函数）：保持 functor 名称
                cpp_type = existing_cpp_type

            node.set_extension_param("cpp_type", cpp_type)
            includes = get_type_includes(type_info, node)
            ir_module.require_includes(includes)

        # 函数节点需要特殊处理返回类型
        if isinstance(node, IRFunctionNode):
            # 对于函数节点，需要处理返回类型
            return_type_to_convert = None

            # 优先从推断的函数类型中获取返回类型
            if type_info and type_info.return_type:
                return_type_to_convert = type_info.return_type
            elif node.return_type:
                # 如果 type_info.return_type 为空，从 node.return_type（Python 侧的类型注解）解析
                from mcpy.type_system import parse_type

                return_type_to_convert = parse_type(node.return_type, node)

            # 转换返回类型为 C++ 类型
            if return_type_to_convert:
                if (
                    node.is_async
                    and return_type_to_convert.kind == TypeKind.GENERIC
                    and return_type_to_convert.name == "Coroutine"
                ):
                    # 协程的返回类型是在第3个参数：Coroutine[Any, Any, T] -> T
                    return_type_to_convert = return_type_to_convert.args[2]

                # 特殊处理：None 作为函数返回类型时映射为 void
                if return_type_to_convert.kind == TypeKind.NONE:
                    return_cpp_type = "void"
                else:
                    return_cpp_type = get_cpp_type(return_type_to_convert, node)

                if return_cpp_type:
                    # async 函数需要包装返回类型为 mc::result<T>
                    if node.is_async:
                        return_cpp_type = f"mc::result<{return_cpp_type}>"
                        ir_module.require_include("<mc/result.h>")

                        # 无 await 的 async 函数需要包装返回值
                        has_await = node.get_extension_param(
                            "has_await_expressions", False
                        )
                        if not has_await:
                            node.set_extension_param("wrap_return_as_result", True)

                    node.set_extension_param("cpp_return_type", return_cpp_type)
                    return_includes = get_type_includes(return_type_to_convert, node)
                    ir_module.require_includes(return_includes)

    def _convert_loop_var_types(self, node: IRNode, ir_module: IRModule):
        """转换循环变量的类型信息

        检查节点是否有循环变量类型信息，并将其转换为 C++ 类型。

        Args:
            node: IR 节点
            ir_module: IR 模块
        """
        # 先收集需要转换的循环变量类型，避免在迭代时修改字典
        loop_var_types = []
        for param_name, param_value in node.extension_params.items():
            if param_name.startswith("loop_var_") and param_name.endswith("_type"):
                # 提取变量名
                var_name = param_name[9:-5]  # 去掉 "loop_var_" 前缀和 "_type" 后缀
                loop_var_types.append((var_name, param_value))

        # 转换收集到的循环变量类型
        for var_name, param_value in loop_var_types:
            if isinstance(param_value, TypeInfo):
                cpp_type = get_cpp_type(param_value, node)
                if cpp_type:
                    # 存储 C++ 类型
                    cpp_type_key = f"loop_var_{var_name}_cpp_type"
                    node.set_extension_param(cpp_type_key, cpp_type)

                    # 收集所需头文件
                    includes = get_type_includes(param_value, node)
                    ir_module.require_includes(includes)

    def _analyze_lambda_capture_strategy(self, node: IRNode, ir_module: IRModule):
        """分析 lambda 闭包捕获策略和使用模式

        决定每个捕获变量的捕获方式：
        - 基础类型（int, float, bool, str）：按值捕获
        - 容器类型（list, dict, set, variant）：按引用捕获（使用 const 引用）
        - 自定义类型：根据大小和可拷贝性决定

        Args:
            node: IR 节点
            ir_module: IR 模块
        """

        if not isinstance(node, IRLambda):
            return

        # 决定使用模式（cpp_lambda vs extension）
        usage_mode = self._determine_lambda_usage_mode(node)
        node.usage_mode = usage_mode

        # 获取已分析的捕获变量类型（由类型推导完成）
        captured_var_types = node.get_extension_param("captured_var_types") or {}

        # 分析每个捕获变量的捕获策略
        capture_strategies = {}
        captured_cpp_types = {}

        for var_name, var_type in captured_var_types.items():
            # 将 TypeInfo 转换为 C++ 类型
            cpp_type = get_cpp_type(var_type, node)
            if not cpp_type:
                cpp_type = "mc::variant"  # 保守策略

            captured_cpp_types[var_name] = cpp_type

            # 获取闭包变量的使用信息（由数据流分析提供）
            closure_usage = node.get_extension_param(f"closure_usage_{var_name}")

            # 决定捕获策略（考虑数据流分析结果）
            strategy = self._determine_capture_strategy_with_dataflow(
                cpp_type, var_type, closure_usage
            )
            capture_strategies[var_name] = strategy

        # 存储分析结果
        node.set_extension_param("captured_cpp_types", captured_cpp_types)
        node.set_extension_param("capture_strategies", capture_strategies)

        # 为所有 lambda 生成稳定的 functor 名称（关键！）
        # 使用模块名 + 源码位置确保跨文件唯一性和生成稳定性
        functor_name = node.get_extension_param("functor_name") or node.functor_name
        if not functor_name:
            # 获取模块名（优先使用 module_path，然后 module_full_name）
            import os

            if hasattr(ir_module, "module_path") and ir_module.module_path:
                # 从文件路径提取文件名（不含扩展名）
                module_name = os.path.splitext(
                    os.path.basename(str(ir_module.module_path))
                )[0]
            elif hasattr(ir_module, "module_full_name") and ir_module.module_full_name:
                module_name = ir_module.module_full_name
            else:
                module_name = "unknown"

            # 清理模块名：移除非法字符
            safe_module_name = (
                module_name.replace("/", "_").replace(".", "_").replace("-", "_")
            )

            # 获取列偏移
            col_offset = node.get_extension_param("col_offset") or 0

            # 生成稳定的 functor 名称：模块名 + 行号 + 列号
            # 这确保了：
            # 1. 跨文件唯一性（模块名不同）
            # 2. 生成稳定性（只要源码位置不变，名称就不变）
            # 3. Git 友好（便于代码审查）
            if hasattr(node, "source_line") and node.source_line > 0:
                functor_name = (
                    f"__lambda_{safe_module_name}_{node.source_line}_{col_offset}"
                )
            else:
                # 降级方案
                functor_name = (
                    f"__lambda_{safe_module_name}_unknown_{id(node) & 0xFFFF}"
                )

            node.functor_name = functor_name
            node.set_extension_param("functor_name", functor_name)

        # 为 lambda 参数设置索引和类型（用于代码生成）
        lambda_type = node.get_extension_param("inferred_type")

        for i, param in enumerate(node.params):
            param.set_extension_param("lambda_arg_index", i)

            # 从 lambda 整体类型中提取参数类型
            if lambda_type and hasattr(lambda_type, "args") and lambda_type.args:
                if i < len(lambda_type.args):
                    param_type = lambda_type.args[i]
                    # 设置参数的推导类型（用于强类型 operator() 生成）
                    param.set_extension_param("inferred_type", param_type)

                    # 同时设置参数的 cpp_type（用于默认值生成）
                    param_cpp_type = get_cpp_type(param_type, node)
                    if param_cpp_type:
                        param.set_extension_param("cpp_type", param_cpp_type)

        # 为栈模式的 lambda 设置 cpp_type
        storage_mode = node.get_extension_param("storage_mode")
        if storage_mode == "stack":
            # 栈模式：cpp_type 是具体的 functor 名称
            node.set_extension_param("cpp_type", functor_name)
        else:
            # 堆模式：cpp_type 是 mc::variant
            node.set_extension_param("cpp_type", "mc::variant")

        # 确保包含必要的头文件
        ir_module.require_include("<mcpy/builtins/lambda.h>")

    def _fix_assignment_target_type(self, node: IRNode, ir_module: IRModule):
        """修复 assignment target 的类型

        为赋值目标变量设置 cpp_type，从右侧表达式传递类型信息。
        这是一个通用的类型传播机制，适用于所有赋值操作：
        - lambda 赋值：x = lambda: ...
        - 变量赋值：y = x
        - 函数调用：z = func()
        - 表达式赋值：w = a + b
        - 字面量赋值：v = [1, 2, 3]
        - 等等...

        Args:
            node: IR 节点
            ir_module: IR 模块
        """
        if not isinstance(node, IRAssign):
            return

        # 获取目标变量
        target = node.get_target_expr()
        if not target or not isinstance(target, IRName):
            return

        # 通用逻辑：从右侧表达式获取 cpp_type
        value_cpp_type = node.value.get_extension_param("cpp_type")

        # 特殊处理：lambda 的 storage_mode（只有 lambda 需要）
        storage_mode = None
        if isinstance(node.value, IRLambda):
            storage_mode = node.value.get_extension_param("storage_mode")

        # 如果右侧表达式没有 cpp_type，尝试从 inferred_type 推导
        if not value_cpp_type:
            from mcpy.type_system import get_cpp_type

            value_inferred_type = node.value.get_extension_param("inferred_type")
            if value_inferred_type:
                value_cpp_type = get_cpp_type(value_inferred_type, node.value)

        if not value_cpp_type:
            return

        # 为赋值节点和目标变量设置 cpp_type（关键！）
        node.set_extension_param("cpp_type", value_cpp_type)
        target.set_extension_param("cpp_type", value_cpp_type)

        if storage_mode == "stack":
            # 栈模式：记录变量名到 functor 名称的映射
            stack_lambda_vars = ir_module.get_extension_param("stack_lambda_vars") or {}
            stack_lambda_vars[target.id] = value_cpp_type
            ir_module.set_extension_param("stack_lambda_vars", stack_lambda_vars)

    def _mark_shared_ptr_variables(
        self, func_node: IRFunctionNode, ir_module: IRModule
    ):
        """标记需要 shared_ptr 包装的变量（函数级别后处理）

        遍历函数中的所有 lambda，收集 shared_ref 捕获的变量，
        然后标记对应的变量声明节点。

        这避免了代码生成阶段重复遍历所有 lambda。

        Args:
            func_node: 函数节点
            ir_module: IR 模块
        """
        # 步骤 1：收集所有需要 shared_ptr 的变量名
        shared_ptr_vars = self._collect_shared_ptr_variables(func_node)

        # 步骤 2：标记变量声明节点
        if shared_ptr_vars:
            self._mark_assignment_nodes_for_shared_ptr(func_node, shared_ptr_vars)

    def _collect_shared_ptr_variables(self, func_node: IRFunctionNode) -> set:
        """收集函数中所有需要 shared_ptr 包装的变量名（支持 Lambda 和嵌套函数）

        Args:
            func_node: 函数节点

        Returns:
            需要 shared_ptr 的变量名集合
        """
        shared_ptr_vars = set()

        # 遍历函数中的所有 lambda 和嵌套函数
        for node in func_node.walk():
            # Lambda 或嵌套函数
            if isinstance(node, IRLambda):
                closure_node = node
            elif isinstance(node, IRFunctionNode) and node != func_node:
                # 嵌套函数（不是当前函数本身）
                if self._is_nested_function_of(node, func_node):
                    closure_node = node
                else:
                    continue
            else:
                continue

            # 获取闭包的捕获策略
            capture_strategies = (
                closure_node.get_extension_param("capture_strategies") or {}
            )

            # 收集 shared_ref 捕获的变量
            for var_name, strategy in capture_strategies.items():
                if strategy == "shared_ref":
                    shared_ptr_vars.add(var_name)

        return shared_ptr_vars

    def _is_nested_function_of(
        self, nested_func: IRFunctionNode, outer_func: IRFunctionNode
    ) -> bool:
        """判断 nested_func 是否是 outer_func 的直接嵌套函数

        Args:
            nested_func: 可能的嵌套函数
            outer_func: 外层函数

        Returns:
            bool: 是否是直接嵌套函数
        """
        parent = nested_func.parent
        while parent:
            if parent == outer_func:
                return True
            if isinstance(parent, IRFunctionNode):
                return False
            parent = parent.parent
        return False

    def _mark_assignment_nodes_for_shared_ptr(
        self, func_node: IRFunctionNode, shared_ptr_vars: set
    ):
        """标记声明了 shared_ptr 变量的赋值节点

        Args:
            func_node: 函数节点
            shared_ptr_vars: 需要 shared_ptr 的变量名集合
        """
        # 遍历函数中的所有赋值节点
        for assign_node in func_node.walk():
            if not isinstance(assign_node, IRAssign):
                continue

            # 检查赋值目标是否包含需要 shared_ptr 的变量
            target_exprs = assign_node.get_target_exprs()
            for target in target_exprs:
                # 只处理简单变量名（IRName）
                if not isinstance(target, IRName):
                    continue

                if target.id in shared_ptr_vars:
                    # 标记赋值节点需要 shared_ptr
                    assign_node.set_extension_param("needs_shared_ptr", True)

                    # 标记具体的变量名（用于多目标赋值场景）
                    marked_vars = assign_node.get_extension_param(
                        "shared_ptr_vars", set()
                    )
                    marked_vars.add(target.id)
                    assign_node.set_extension_param("shared_ptr_vars", marked_vars)

    def _fix_stack_lambda_references(self, node: IRNode, ir_module: IRModule):
        """修复所有引用闭包（lambda 或嵌套函数）变量的 IRName 节点的 cpp_type

        为所有闭包变量的引用设置正确的 cpp_type：
        - 栈模式：functor 名称（如 __lambda_xxx 或 __nested_func_xxx）
        - 堆模式：mc::variant


        Args:
            node: IR 节点
            ir_module: IR 模块
        """
        if not isinstance(node, IRName):
            return

        # 从 ir_module 获取栈闭包变量映射
        stack_lambda_vars = ir_module.get_extension_param("stack_lambda_vars") or {}

        # 检查当前变量是否是栈闭包
        # 重要：不能仅凭变量名判断，必须查找符号确认
        if node.id in stack_lambda_vars:
            # 查找符号，确认是否真的是栈闭包变量
            symbol = node.lookup_symbol(node.id)
            if symbol and symbol.ir_node:
                # 情况1：符号指向 IRFunctionNode（嵌套函数）
                if isinstance(symbol.ir_node, IRFunctionNode):
                    # 确认是栈嵌套函数
                    storage_mode = symbol.ir_node.get_extension_param("storage_mode")
                    if storage_mode == "stack":
                        functor_name = stack_lambda_vars[node.id]
                        node.set_extension_param(
                            "cpp_type", functor_name, replace_existing=True
                        )
                        node.set_extension_param("is_stack_lambda_ref", True)
                        return

                # 情况2：符号指向 IRName（lambda 赋值）
                if isinstance(symbol.ir_node, IRName):
                    # 符号指向的是 IRName，查找其父节点（应该是 IRAssign）
                    parent = symbol.ir_node.parent
                    if isinstance(parent, IRAssign):
                        # 检查赋值的值是 lambda 还是嵌套函数
                        if isinstance(parent.value, IRLambda):
                            # 确认是栈 lambda
                            storage_mode = parent.value.get_extension_param(
                                "storage_mode"
                            )
                            if storage_mode == "stack":
                                functor_name = stack_lambda_vars[node.id]
                                node.set_extension_param(
                                    "cpp_type", functor_name, replace_existing=True
                                )
                                node.set_extension_param("is_stack_lambda_ref", True)
                                return
                        elif isinstance(parent.value, IRFunctionNode):
                            # 确认是栈嵌套函数（赋值形式）
                            storage_mode = parent.value.get_extension_param(
                                "storage_mode"
                            )
                            if storage_mode == "stack":
                                functor_name = stack_lambda_vars[node.id]
                                node.set_extension_param(
                                    "cpp_type", functor_name, replace_existing=True
                                )
                                node.set_extension_param("is_stack_lambda_ref", True)
                                return

        # 检查当前变量是否是堆模式 lambda
        symbol = node.lookup_symbol(node.id)
        if symbol and symbol.ir_node and isinstance(symbol.ir_node, IRName):
            # 符号指向的是 IRName，查找其父节点（应该是 IRAssign）
            parent = symbol.ir_node.parent
            if isinstance(parent, IRAssign) and isinstance(parent.value, IRLambda):
                # 获取 lambda 的 cpp_type
                lambda_cpp_type = parent.value.get_extension_param("cpp_type")
                if lambda_cpp_type:
                    # 设置到当前引用节点
                    node.set_extension_param(
                        "cpp_type", lambda_cpp_type, replace_existing=True
                    )

    def _determine_capture_strategy(self, cpp_type: str, type_info: TypeInfo) -> str:
        """确定捕获策略

        Args:
            cpp_type: C++ 类型字符串
            type_info: 类型信息

        Returns:
            捕获策略：
            - "value": 按值捕获（基础类型）
            - "const_ref": 按 const 引用捕获（容器类型）
            - "shared_ref": 使用智能指针捕获（需要共享所有权的类型）
        """
        # 优先使用类型信息中的 pass_by 配置
        if type_info.pass_by:
            # 映射 pass_by 到捕获策略
            # pass_by 可能的值：'value', 'const_ref', 'rvalue_ref', 'auto'
            if type_info.pass_by == "value":
                return "value"  # 按值捕获
            elif type_info.pass_by in ("const_ref", "auto"):
                return "const_ref"  # 按常量引用捕获
            # rvalue_ref 等特殊情况使用 const_ref（保守策略）
            else:
                return "const_ref"

        # 回退策略：基于 C++ 类型字符串判断（用于没有配置的类型）
        # variant 类型：按常量引用捕获
        if cpp_type == "mc::variant":
            return "const_ref"

        # 智能指针类型：按值捕获（智能指针本身轻量）
        if cpp_type.startswith("std::shared_ptr<") or cpp_type.startswith(
            "std::unique_ptr<"
        ):
            return "value"

        # 未知类型：保守策略，使用 const 引用
        return "const_ref"

    def _determine_capture_strategy_with_dataflow(
        self, cpp_type: str, type_info: TypeInfo, closure_usage
    ) -> str:
        """确定捕获策略

        Args:
            cpp_type: C++ 类型字符串
            type_info: 类型信息
            closure_usage: 闭包变量使用信息（ClosureVariableUsage 或 None）

        Returns:
            捕获策略：
            - "value": 按值捕获（基础类型）
            - "const_ref": 按 const 引用捕获（容器类型）
            - "shared_ref": 使用智能指针捕获（需要共享所有权）
        """
        # 如果有数据流分析结果，检查是否需要使用 shared_ref
        if closure_usage:
            # 特殊处理：参数（CLOSURE_VARIABLE）被多个闭包共享不需要 shared_ref
            # 参数的生命周期由外层函数保证，不需要智能指针
            is_parameter = (
                closure_usage.is_parameter
                if hasattr(closure_usage, "is_parameter")
                else False
            )

            # 如果变量被多个 lambda 共享，或者在 lambda 定义后被修改，使用 shared_ref
            # 但参数除外
            if not is_parameter and (
                closure_usage.shared_by_multiple_lambdas
                or closure_usage.modified_after_capture
            ):
                return "shared_ref"

        # 否则使用基于类型的默认策略
        return self._determine_capture_strategy(cpp_type, type_info)

    def _determine_lambda_usage_mode(self, node: IRLambda) -> str:
        """决定 Lambda 使用模式

        决策规则：
        - 已设置usage_mode="cpp_lambda" → 保持（如async then回调）
        - 立即调用场景 → cpp_lambda（零开销）
        - 其他场景 → extension（可存储）

        Args:
            node: Lambda 节点

        Returns:
            使用模式："cpp_lambda" 或 "extension"
        """
        # 如果已经设置为cpp_lambda（如async then回调），保持不变
        if node.usage_mode == "cpp_lambda":
            return "cpp_lambda"

        # 检查是否是async then回调（AsyncTransformPass设置）
        if node.get_extension_param("async_then_callback"):
            return "cpp_lambda"

        # 检查是否是立即调用场景：(lambda x: x + 1)(5)
        parent = node.parent
        if isinstance(parent, IRCall) and parent.func == node:
            # 立即调用：使用零开销 C++ lambda
            return "cpp_lambda"

        # 其他场景：使用 extension（可存储、可传递）
        return "extension"

    def _analyze_nested_function_capture_strategy(
        self, node: IRNode, ir_module: IRModule
    ):
        """分析嵌套函数闭包捕获策略（复用 Lambda 逻辑 100%）

        Args:
            node: IR 节点
            ir_module: IR 模块
        """
        if not isinstance(node, IRFunctionNode):
            return

        if not self._is_nested_function(node):
            return

        # 获取已分析的捕获变量类型（由类型推导完成）
        captured_var_types = node.get_extension_param("captured_var_types") or {}

        # 分析每个捕获变量的捕获策略
        capture_strategies = {}
        captured_cpp_types = {}

        for var_name, var_type in captured_var_types.items():
            # 将 TypeInfo 转换为 C++ 类型
            cpp_type = get_cpp_type(var_type, node)
            if not cpp_type:
                cpp_type = "mc::variant"  # 保守策略

            captured_cpp_types[var_name] = cpp_type

            # 获取闭包变量的使用信息（由数据流分析提供）
            closure_usage = node.get_extension_param(f"closure_usage_{var_name}")

            # 决定捕获策略（完全复用 Lambda 逻辑）
            strategy = self._determine_capture_strategy_with_dataflow(
                cpp_type, var_type, closure_usage
            )
            capture_strategies[var_name] = strategy

        # 存储分析结果
        node.set_extension_param("captured_cpp_types", captured_cpp_types)
        node.set_extension_param("capture_strategies", capture_strategies)

        # 生成稳定的 functor 名称
        functor_name = node.get_extension_param("functor_name")
        if not functor_name:
            import os

            # 获取模块名
            if hasattr(ir_module, "module_path") and ir_module.module_path:
                module_name = os.path.splitext(
                    os.path.basename(str(ir_module.module_path))
                )[0]
            elif hasattr(ir_module, "module_full_name") and ir_module.module_full_name:
                module_name = ir_module.module_full_name
            else:
                module_name = "unknown"

            safe_module_name = (
                module_name.replace("/", "_").replace(".", "_").replace("-", "_")
            )

            # 查找外层函数名（与 NestedFunctionEmitter 保持一致）
            parent = node.parent
            while parent and parent.node_type != "function":
                parent = parent.parent
            outer_func_name = (
                parent.name if parent and parent.node_type == "function" else "func"
            )

            # 生成稳定的 functor 名称：使用模块名 + 外层函数名 + 函数名 + 行号
            if hasattr(node, "source_line") and node.source_line > 0:
                functor_name = f"__nested_func_{safe_module_name}_{outer_func_name}_{node.name}_{node.source_line}"
            else:
                functor_name = f"__nested_func_{safe_module_name}_{outer_func_name}_{node.name}_{id(node) & 0xFFFF}"

            node.set_extension_param("functor_name", functor_name)

        # 决定存储模式
        storage_mode = self._determine_nested_function_storage_mode(node, ir_module)
        node.set_extension_param("storage_mode", storage_mode)

        # 为栈模式的嵌套函数设置 cpp_type
        # 注意：functor_name 现在已经存储在 node 的 extension_params 中
        final_functor_name = node.get_extension_param("functor_name")
        if storage_mode == "stack":
            node.set_extension_param("cpp_type", final_functor_name)

            # 记录栈嵌套函数到 stack_lambda_vars（供 _fix_stack_lambda_references 使用）
            # 嵌套函数定义在 Python 中是变量赋值
            stack_lambda_vars = ir_module.get_extension_param("stack_lambda_vars") or {}
            stack_lambda_vars[node.name] = final_functor_name
            ir_module.set_extension_param(
                "stack_lambda_vars", stack_lambda_vars, replace_existing=True
            )
        else:
            node.set_extension_param("cpp_type", "mc::variant")

        # 设置 cpp_return_type（函数节点需要）
        # 嵌套函数的返回类型推导（与 _convert_node_type 中的逻辑一致）
        func_type = node.get_extension_param("inferred_type")
        if func_type and func_type.return_type:
            # 特殊处理：如果返回类型是函数类型，强制使用 mc::variant
            # 函数作为值（返回、赋值、传递）时始终是 mc::variant
            if func_type.return_type.kind == TypeKind.FUNCTION:
                node.set_extension_param("", "mc::variant")
            else:
                # 特殊处理：None 作为函数返回类型时映射为 void
                if func_type.return_type.kind == TypeKind.NONE:
                    return_cpp_type = "void"
                else:
                    # 获取返回类型的 C++ 类型
                    return_cpp_type = get_cpp_type(func_type.return_type, node)

                if return_cpp_type:
                    # async 函数需要包装返回类型为 mc::result<T>
                    if node.is_async:
                        return_cpp_type = f"mc::result<{return_cpp_type}>"
                        ir_module.require_include("<mc/result.h>")

                        # 无 await 的 async 函数需要包装返回值
                        has_await = node.get_extension_param(
                            "has_await_expressions", False
                        )
                        if not has_await:
                            node.set_extension_param("wrap_return_as_result", True)

                    node.set_extension_param("cpp_return_type", return_cpp_type)

        # 确保包含必要的头文件
        ir_module.require_include("<mcpy/builtins/lambda.h>")

    def _is_nested_function(self, func_node: IRFunctionNode) -> bool:
        """判断是否为嵌套函数

        Args:
            func_node: 函数节点

        Returns:
            bool: 是否为嵌套函数
        """
        parent = func_node.parent
        while parent:
            if isinstance(parent, (IRFunctionNode, IRMethodNode)):
                return True
            parent = parent.parent
        return False

    def _determine_nested_function_storage_mode(
        self, node: IRFunctionNode, ir_module: IRModule
    ) -> str:
        """决定嵌套函数的存储模式

        Args:
            node: 函数节点
            ir_module: IR 模块

        Returns:
            存储模式："stack" 或 "heap"
        """
        # 简化规则：检查嵌套函数名是否在外层函数的 return 语句中被使用
        containing_func = self._get_containing_function(node)
        if containing_func:
            # 遍历所有 return 语句，检查是否返回了嵌套函数
            for ret_node in containing_func.walk():
                if isinstance(ret_node, IRReturn) and ret_node.expr:
                    # 检查 return 值是否包含嵌套函数名
                    ret_value_names = set()
                    for name_node in ret_node.expr.walk():
                        if isinstance(name_node, IRName):
                            ret_value_names.add(name_node.id)

                    if node.name in ret_value_names:
                        # 嵌套函数被返回 → 堆分配
                        return "heap"

        # 默认栈分配（局部使用）
        return "stack"

    def _get_containing_function(self, node: IRNode) -> IRFunctionNode | None:
        """获取包含当前节点的函数

        Args:
            node: IR 节点

        Returns:
            包含函数节点，如果没有则返回 None
        """
        parent = node.parent
        while parent:
            if isinstance(parent, (IRFunctionNode, IRMethodNode)):
                return parent
            parent = parent.parent
        return None

    def _convert_async_state_vars(self, func_node: IRFunctionNode) -> None:
        """转换异步状态变量的类型（Python TypeInfo -> C++ 类型字符串）

        这个方法处理 async_state_vars 字典中的类型信息，
        将 Python TypeInfo 对象或 VariablePromotionInfo 对象转换为 C++ 类型字符串，供代码生成器使用。

        Args:
            func_node: 异步函数节点
        """
        state_vars = func_node.get_extension_param("async_state_vars")
        if not state_vars:
            return

        cpp_state_vars = {}
        for var_name, value in state_vars.items():
            # 如果已经是字符串，直接使用
            if isinstance(value, str):
                cpp_state_vars[var_name] = value
                continue

            # 如果是 None，使用默认类型
            if value is None:
                cpp_state_vars[var_name] = "mc::variant"
                continue

            # 如果是 VariablePromotionInfo 对象，从中提取 TypeInfo
            type_info = None
            if hasattr(value, "type_info"):
                # VariablePromotionInfo 对象
                type_info = value.type_info
            else:
                # 直接是 TypeInfo 对象
                type_info = value

            # 如果没有类型信息，使用默认类型
            if type_info is None:
                cpp_state_vars[var_name] = "mc::variant"
                continue

            # 转换 TypeInfo 为 C++ 类型
            cpp_type = get_cpp_type(type_info, func_node)
            cpp_state_vars[var_name] = cpp_type if cpp_type else "mc::variant"

        # 更新为 C++ 类型（必须使用 replace_existing=True 来覆盖之前的 TypeInfo）
        func_node.set_extension_param(
            "async_state_vars", cpp_state_vars, replace_existing=True
        )
